from fastapi import APIRouter, WebSocket, WebSocketDisconnect, Depends, Query
from typing import Dict, List
import json
import asyncio
import logging
from datetime import datetime

from app.dependencies import get_current_user_ws, get_db_session
from app.models.database import User
from sqlalchemy.ext.asyncio import AsyncSession

router = APIRouter()
logger = logging.getLogger(__name__)

class ConnectionManager:
    """WebSocket连接管理器"""
    
    def __init__(self):
        # 存储活跃连接 {user_id: {connection_id: websocket}}
        self.active_connections: Dict[str, Dict[str, WebSocket]] = {}
        # 任务订阅 {task_id: [user_ids]}
        self.task_subscriptions: Dict[str, List[str]] = {}
        # 系统通知订阅 {user_id: websocket}
        self.notification_subscribers: Dict[str, WebSocket] = {}
    
    async def connect(self, websocket: WebSocket, user_id: str, connection_id: str):
        """建立WebSocket连接"""
        await websocket.accept()
        
        if user_id not in self.active_connections:
            self.active_connections[user_id] = {}
        
        self.active_connections[user_id][connection_id] = websocket
        logger.info(f"WebSocket connected: user={user_id}, connection={connection_id}")
    
    def disconnect(self, user_id: str, connection_id: str):
        """断开WebSocket连接"""
        if user_id in self.active_connections:
            self.active_connections[user_id].pop(connection_id, None)
            if not self.active_connections[user_id]:
                del self.active_connections[user_id]
        
        # 清理任务订阅
        for task_id, subscribers in self.task_subscriptions.items():
            if user_id in subscribers:
                subscribers.remove(user_id)
        
        # 清理通知订阅
        self.notification_subscribers.pop(user_id, None)
        
        logger.info(f"WebSocket disconnected: user={user_id}, connection={connection_id}")
    
    async def send_personal_message(self, message: dict, user_id: str):
        """发送个人消息"""
        if user_id in self.active_connections:
            disconnected = []
            for connection_id, websocket in self.active_connections[user_id].items():
                try:
                    await websocket.send_text(json.dumps(message))
                except Exception as e:
                    logger.error(f"Failed to send message to {user_id}/{connection_id}: {e}")
                    disconnected.append(connection_id)
            
            # 清理断开的连接
            for connection_id in disconnected:
                self.disconnect(user_id, connection_id)
    
    async def send_task_update(self, task_id: str, message: dict):
        """发送任务更新消息"""
        if task_id in self.task_subscriptions:
            for user_id in self.task_subscriptions[task_id]:
                await self.send_personal_message(message, user_id)
    
    async def broadcast_notification(self, message: dict):
        """广播系统通知"""
        disconnected = []
        for user_id, websocket in self.notification_subscribers.items():
            try:
                await websocket.send_text(json.dumps(message))
            except Exception as e:
                logger.error(f"Failed to send notification to {user_id}: {e}")
                disconnected.append(user_id)
        
        # 清理断开的连接
        for user_id in disconnected:
            self.notification_subscribers.pop(user_id, None)
    
    def subscribe_task(self, task_id: str, user_id: str):
        """订阅任务更新"""
        if task_id not in self.task_subscriptions:
            self.task_subscriptions[task_id] = []
        
        if user_id not in self.task_subscriptions[task_id]:
            self.task_subscriptions[task_id].append(user_id)
    
    def unsubscribe_task(self, task_id: str, user_id: str):
        """取消订阅任务更新"""
        if task_id in self.task_subscriptions:
            if user_id in self.task_subscriptions[task_id]:
                self.task_subscriptions[task_id].remove(user_id)
    
    def subscribe_notifications(self, user_id: str, websocket: WebSocket):
        """订阅系统通知"""
        self.notification_subscribers[user_id] = websocket
    
    def get_connection_count(self) -> int:
        """获取连接总数"""
        total = 0
        for user_connections in self.active_connections.values():
            total += len(user_connections)
        return total

# 全局连接管理器
manager = ConnectionManager()

@router.websocket("/tasks/{task_id}/status")
async def websocket_task_status(
    websocket: WebSocket,
    task_id: str,
    token: str = Query(..., description="认证令牌"),
    db: AsyncSession = Depends(get_db_session)
):
    """任务状态实时推送"""
    try:
        # WebSocket认证
        user = await get_current_user_ws(websocket, token, db)
        connection_id = f"task_status_{task_id}_{user.id}"
        
        await manager.connect(websocket, user.id, connection_id)
        manager.subscribe_task(task_id, user.id)
        
        # 发送初始状态
        await websocket.send_text(json.dumps({
            "type": "connected",
            "message": f"已连接到任务 {task_id} 的状态推送",
            "timestamp": datetime.now().isoformat()
        }))
        
        try:
            while True:
                # 保持连接活跃，接收客户端心跳
                data = await websocket.receive_text()
                
                # 处理客户端消息
                try:
                    message = json.loads(data)
                    if message.get("type") == "ping":
                        await websocket.send_text(json.dumps({
                            "type": "pong",
                            "timestamp": datetime.now().isoformat()
                        }))
                except json.JSONDecodeError:
                    pass
                
        except WebSocketDisconnect:
            manager.disconnect(user.id, connection_id)
            manager.unsubscribe_task(task_id, user.id)
            
    except Exception as e:
        logger.error(f"WebSocket task status error: {e}")
        try:
            await websocket.close()
        except:
            pass

@router.websocket("/tasks/{task_id}/logs")
async def websocket_task_logs(
    websocket: WebSocket,
    task_id: str,
    token: str = Query(..., description="认证令牌"),
    db: AsyncSession = Depends(get_db_session)
):
    """任务日志实时推送"""
    try:
        # WebSocket认证
        user = await get_current_user_ws(websocket, token, db)
        connection_id = f"task_logs_{task_id}_{user.id}"
        
        await manager.connect(websocket, user.id, connection_id)
        
        # 发送初始消息
        await websocket.send_text(json.dumps({
            "type": "connected",
            "message": f"已连接到任务 {task_id} 的日志推送",
            "timestamp": datetime.now().isoformat()
        }))
        
        try:
            while True:
                data = await websocket.receive_text()
                
                # 处理客户端消息
                try:
                    message = json.loads(data)
                    if message.get("type") == "ping":
                        await websocket.send_text(json.dumps({
                            "type": "pong",
                            "timestamp": datetime.now().isoformat()
                        }))
                except json.JSONDecodeError:
                    pass
                
        except WebSocketDisconnect:
            manager.disconnect(user.id, connection_id)
            
    except Exception as e:
        logger.error(f"WebSocket task logs error: {e}")
        try:
            await websocket.close()
        except:
            pass

@router.websocket("/notifications")
async def websocket_notifications(
    websocket: WebSocket,
    token: str = Query(..., description="认证令牌"),
    db: AsyncSession = Depends(get_db_session)
):
    """系统通知实时推送"""
    try:
        # WebSocket认证
        user = await get_current_user_ws(websocket, token, db)
        connection_id = f"notifications_{user.id}"
        
        await manager.connect(websocket, user.id, connection_id)
        manager.subscribe_notifications(user.id, websocket)
        
        # 发送初始消息
        await websocket.send_text(json.dumps({
            "type": "connected",
            "message": "已连接到系统通知推送",
            "timestamp": datetime.now().isoformat()
        }))
        
        try:
            while True:
                data = await websocket.receive_text()
                
                # 处理客户端消息
                try:
                    message = json.loads(data)
                    if message.get("type") == "ping":
                        await websocket.send_text(json.dumps({
                            "type": "pong",
                            "timestamp": datetime.now().isoformat()
                        }))
                except json.JSONDecodeError:
                    pass
                
        except WebSocketDisconnect:
            manager.disconnect(user.id, connection_id)
            
    except Exception as e:
        logger.error(f"WebSocket notifications error: {e}")
        try:
            await websocket.close()
        except:
            pass

# 辅助函数：发送任务状态更新
async def broadcast_task_status(task_id: str, status_data: dict):
    """广播任务状态更新"""
    message = {
        "type": "status_update",
        "data": {
            "task_id": task_id,
            "timestamp": datetime.now().isoformat(),
            **status_data
        }
    }
    await manager.send_task_update(task_id, message)

# 辅助函数：发送任务日志
async def broadcast_task_log(task_id: str, log_data: dict):
    """广播任务日志"""
    message = {
        "type": "log_entry",
        "data": {
            "task_id": task_id,
            "timestamp": datetime.now().isoformat(),
            **log_data
        }
    }
    await manager.send_task_update(task_id, message)

# 辅助函数：发送系统通知
async def broadcast_system_notification(title: str, message: str, level: str = "info"):
    """广播系统通知"""
    notification = {
        "type": "notification",
        "data": {
            "id": f"notif_{datetime.now().timestamp()}",
            "title": title,
            "message": message,
            "level": level,
            "timestamp": datetime.now().isoformat()
        }
    }
    await manager.broadcast_notification(notification)

# 获取连接统计信息
def get_websocket_stats() -> dict:
    """获取WebSocket连接统计"""
    return {
        "total_connections": manager.get_connection_count(),
        "active_users": len(manager.active_connections),
        "task_subscriptions": len(manager.task_subscriptions),
        "notification_subscribers": len(manager.notification_subscribers)
    }
